'''
Description:          
Author: WANG Hao
Date: 2022-09-15 11:45:46
LastEditTime: 2022-09-15 11:45:46
LastEditors: WANG Hao
'''


#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@Copyright:  Jihua Lab 2022.
@File    :   sea_dual_arm.py    
@Author :   HaoWANG, Foshan，China
@Contact :   wanghao@jihualab.com
@License :   JHL ( not BSD)

@Description:  



@Create Time   :   2022/9/15 0008  9:19
@Modify Time      @Author    @Version    @Desciption
------------      -------    --------    -----------
2022/9/15 0008 9:19   wanghao      1.1.0         None
20220915 
20220914 TODO: Three Threads: 1 - Dual arm motion control 
                2 - records the left pitch joint angle state 
                3- records the right pitch joint angle state 

'''

# import ADS1263 lib
from concurrent.futures import thread
import time
import ADS1263
import queue

# import dual arm motion control lib

import logging
import threading

import uservo
import serial

from MotionGroup import MotionGroup
from ServoPose import Pose
from wrist_gripper import WristGripperControl
from robot_arm import ArmMotionControl

"""
controller paramters
"""
# 标志
exitFlag = 0
recordFlag = 0

# Global const para
THREAD_MODE_CONTROL = ["control", "Control", "CONTROL"]
THREAD_MODE_RECORD = ["record", "Record", "Records", "RECORD", "RECORDS"]
THREAD_MODE_RECORD_SEA = ["record", "Record", "Records", "RECORD", "RECORDS"]


SHORT_TIME_DELAY = 0.02  # time delay para
MEDIA_TIME_DELAY = 0.10
LONG_TIME_DELAY = 1.0

# recording
SAMPLING_FREQUENCY = 10  # units: Hz
fileName = "robot_arm_joint_records.txt"

# servo motor velocity： mean rotation speed （deg/sec）  
VEL = {'VERY_LOW_SPEED': 10, 'LOW_SPEED': 15, 'MID_SPEED': 30, 'HIGH_SPEED': 60}
# Control period : total time used for one servo motor control cmd,( ms) 
CONTROL_PERIOD = {"SEC_5":5000, "SEC_3":3000, "SEC_2":2000, "SEC_1":1000, "SEC_0.5":500}
CONTROL_MODE = {'Teach': True, 'Repeat': False}




# left arm + left wrist
# right arm + right wrist + left and right gripper angle
"""
INIT POSE - INIT joint state "I_"
"""
DUAL_ARM_INIT_POSE = [-151.0, 138.0, 110.0, 0, 
                        0.0, 23.0, 72.0, 0, 
                        30, 30]
"""
START POSE - TEST START joint state "I"
"""
DUAL_ARM_START_POSE = [-151.0, 138.0, -26.0, 0, 
                        0.0, 23.0, -60.0, 0, 
                        30, 30]
# # 左臂抬升二关节抬升至中间位置
# DUAL_ARM_PRO_POSE_LEFT = [140.0, -30.0, -30.0, 85,
#                           -1, -84, -40, 90, 60, 60]
# # 左臂第一关节向左摆动90度，左臂第二关节俯仰角保持不变
# DUAL_ARM_TARGET_POSE_LEFT = [50.0, -30.0, -30.0, 85,
#                              -1, -84, -40, 90, 60, 60]

# right arm 
# shoulder pithch joint rotation angle (delta theta form start pose)
DUAL_ARM_TARGET_POSE_RIGHT = [-151.0, 138.0, -26.0, 0, 
                        0.0, 23.0+delta_theta_right, -60.0, 0, 
                        30, 30]
delta_theta_right  = 30 

"""
END OF controller paramters
"""


class dfrThread(threading.Thread):
    def __init__(self, threadID, name, controlPeriod, samplingFrequency, dual_arm_obj, sea_serial = 0):
        threading.Thread.__init__(self)
        self.threadID = threadID
        self.threadName = name
        self.controlPeriod = controlPeriod
        self.samplingFrequency = samplingFrequency
        self.dual_arm_obj = dual_arm_obj
        self.sea_serial = sea_serial

    def run(self):
        print("开始线程：" + self.threadName)
        if (self.threadName in THREAD_MODE_CONTROL):
            logging.info("%s: %s" % (self.threadName, time.ctime(time.time())))

            robot_arm_control(self.threadName, self.dual_arm_obj, self.controlPeriod)

        elif self.threadName in THREAD_MODE_RECORD:
            record_servo_info(self.threadName, self.dual_arm_obj, samplingFrequency=SAMPLING_FREQUENCY)
        
        elif self.threadName in THREAD_MODE_RECORD_SEA:
            
            record_sea_info(threadName=self.threadName,sea_serial=self.sea_serial)

        else:
            logging.error("self.threadName should be in THREAD_MODE_RECORD or THREAD_MODE_CONTROL !")

        print("退出线程：" + self.name)


"""
Description：
1. robot arm motion control module for dual-arm SEA joint puse reaction experiments
2. threads for motion control and data records function 
3. PARAMTERS tune for CONTROL PERIOD and VEL\ACC\DEC time  
"""
def robot_arm_control(threadName, dual_arm, controlPeriod=3000, vel = VEL.get('VERY_LOW_SPEED')):
    # dual arm motion control threading
    dual_arm_servo = dual_arm.get_arm_servo_list()
    record_serial = 1

    while dual_arm and (threadName in THREAD_MODE_CONTROL):
        keyInput = input('waitKey to continue input (c or x):')
        if keyInput == 'x':
            dual_arm.set_all_servo_mode(CONTROL_MODE="CIRCLE")
            logging.warning('Interrupt:  The waitKey input (c or x): {}'.format(keyInput))
            break
        elif keyInput == 'c':
            # dual_arm再现
            # move the left and right robot arm to start pose for later trajectory control
            for SERVO_ID in dual_arm_servo:
                dual_arm.set_single_servo_pose(
                    Pose=generate_pose_vtc(angle=DUAL_ARM_START_POSE[SERVO_ID], vel=vel,
                                           iterval=controlPeriod,
                                           t_acc=45,
                                           t_dec=40, power=10000, mean_dps=30, servo_id=SERVO_ID))
            time.sleep(1 * LONG_TIME_DELAY)


            while True:
                loadingKeyInput = input('waitKey to control robot arm and read joint state -> input (y or n):')
                if loadingKeyInput == 'n':
                    # 修改标志位，线程上锁
                    # threadLock.acquire()
                    break
                elif loadingKeyInput == 'y':
                    # threadLock.acquire()
                    recordFlag = 1  # record 0->1
                    exitFlag = 0  # exit 0->0
                    # threadLock.release()
                    
                    #TODO: records shoulder pitch joint servo motor angle in DEGREE 
                    # file_joint_angle = "file_joint_angle_records"+record_serial+".txt"
                    # with open(file_joint_angle, 'w') as f:
                    #     f.open

                    for i in range(5):

                        # target pose trajectory control for left arm servo id = 0,1,2
                        for SERVO_ID in dual_arm_servo:
                            dual_arm.set_single_servo_pose(
                                Pose=generate_pose_vtc(angle=DUAL_ARM_TARGET_POSE_RIGHT[SERVO_ID],
                                                       vel=vel,
                                                       iterval=controlPeriod,
                                                       t_acc=45,
                                                       t_dec=40, power=10000, mean_dps=30, servo_id=SERVO_ID))
                        time.sleep(2 * LONG_TIME_DELAY)
                        # start pose trajectory control
                        for SERVO_ID in dual_arm_servo:
                            dual_arm.set_single_servo_pose(
                                Pose=generate_pose_vtc(angle=DUAL_ARM_START_POSE[SERVO_ID],
                                                       vel=vel,
                                                       iterval=controlPeriod,
                                                       t_acc=45,
                                                       t_dec=40, power=10000, mean_dps=30, servo_id=SERVO_ID))
                        time.sleep(2 * LONG_TIME_DELAY)


                    # target pose trajectory control for right arm servo id = 0,1,2,4,5,6
                    time.sleep(5)
        else:
            logging.warning('waitKey input (c or x): {}'.format(keyInput))


# 1. servo motor real rotation state records\
# 2. record robot arm servo info and SEA joint angle state infomation
def record_servo_info(threadName, dual_arm, samplingFrequency, frame_num = 500):
    while dual_arm and samplingFrequency and (threadName in THREAD_MODE_RECORD):
        loadingKeyInput = input('waitKey to record continue input (y or n):')
        if loadingKeyInput == 'n':
            threadName.exit()
            logging.info("exitFlag = True: threadName.exit() !")
            break

        elif loadingKeyInput == 'y':
            with open(fileName, 'w') as f:
                for frame in range(frame_num):
                    # data file write as txt
                    f.write("frame ID ： {}   {}   {}   {} ".format(frame,
                                                            time.ctime(time.time()),
                                                            dual_arm.read_single_servo_angle(SERVO_ID=1) -
                                                            DUAL_ARM_START_POSE[1],
                                                            dual_arm.read_single_servo_angle(SERVO_ID=5) -
                                                            DUAL_ARM_START_POSE[5]) + '\r')
                    time.sleep(1 / samplingFrequency)
            break
        # logging.info("%s: %s" % (threadName, time.ctime(time.time())))
        # recoding the angle data and the power data


# record robot arm servo info and SEA joint angle state infomation
def record_sea_info(threadName, 
                    sea_serial, 
                    samplingFrequency = 20, 
                    auto_start = True):

    DefaultSEAChannelSer = [0 , 1]
  
    REF = 5.08  # Modify according to actual voltage
    # external AVDD and AVSS(Default), or internal 2.5V
    #servo_list = [0, 1, 4, 5]  # 0,1 left arm ;4,5 right arm shoulder joint yaw and pitch, respectively
    #SEA_Balanced_Angle = []  # static information of sea joint balanced angle
    # def average size

    kernel_size = 3
    frequency = samplingFrequency   # default 20 Hz
    moving_average_a = MovingAverage(kernel_size)
    moving_average_b = MovingAverage(kernel_size)

    current_pitch_right_buffer = []  # channel 1
    current_pitch_left_buffer = []  # channel 0

    if auto_start and sea_serial in DefaultSEAChannelSer:     
        # sea_serial - > channel 0, 1, and 0 for left shoulder pitchjoint
        #                 channel 1 for right shoulder pitch joint SEA state estimation results record 
        if sea_serial == 0:  # channel 0
            try:
                ADC = ADS1263.ADS1263()
                if (ADC.ADS1263_init_ADC1('ADS1263_7200SPS') == -1):
                    exit()
                ADC.ADS1263_SetMode(0)

                # ADC.ADS1263_DAC_Test(1, 1)      # Open IN6
                # ADC.ADS1263_DAC_Test(0, 1)      # Open IN7

                if (ADC.ADS1263_init_ADC2('ADS1263_ADC2_400SPS') == -1):
                    exit()
                while (1):
                    ADC_Value = ADC.ADS1263_GetAll_ADC2()  # get ADC2 value
                    i = sea_serial # IN0 ADC_IN0 Value
                    if (ADC_Value[i] >> 23 == 1):
                        val = REF * 2 - ADC_Value[i] * REF / 0x800000
                        avg = moving_average_a.next(val)
                        logging.info("ADC2 IN%d = -%lf" % (i, (REF * 2 - ADC_Value[i] * REF / 0x800000)))
                        logging.info("     AVG a : ADC2 IN%d = -%lf" % (i, avg))
                    else:
                        val = ADC_Value[i] * REF / 0x7fffff
                        avg = moving_average_a.next(val)
                        logging.info("ADC2 IN%d = %lf" % (i, (ADC_Value[i] * REF / 0x7fffff)))  # 24bit
                        logging.info("     AVG a : ADC2 IN%d = %lf" % (i, avg))
                    angle = (360.0 * avg) / REF  # degree of sea

                    current_pitch_left_buffer.append(angle)
                    logging.info("SEA  ID : {}   - ANGLE : ".format(i)+ format(angle, '.3f'))

                    time.sleep(1 / frequency)

            except IOError as e:
                logging.error(e)

            except KeyboardInterrupt:
                logging.info("ctrl + c: Program end")
                # logging.info(current_yaw_right_buffer,current_pitch_right_buffer)
                logging.info("sea_serial NOT in DefaultSEAChannelSer: threadName.exit() !")
                ADC.ADS1263_Exit()
                threadName.exit()
                exit()


        elif sea_serial ==1:  # channel 1
            try:
                ADC = ADS1263.ADS1263()
                if (ADC.ADS1263_init_ADC1('ADS1263_7200SPS') == -1):
                    exit()
                ADC.ADS1263_SetMode(0)

                # ADC.ADS1263_DAC_Test(1, 1)      # Open IN6
                # ADC.ADS1263_DAC_Test(0, 1)      # Open IN7

                if (ADC.ADS1263_init_ADC2('ADS1263_ADC2_400SPS') == -1):
                    exit()
                while (1):
                    ADC_Value = ADC.ADS1263_GetAll_ADC2()  # get ADC2 value
                    i = sea_serial # IN0 ADC_IN0 Value
                    if (ADC_Value[i] >> 23 == 1):
                        val_b = REF * 2 - ADC_Value[i] * REF / 0x800000
                        avg_b = moving_average_b.next(val_b)
                        logging.info("ADC2 IN%d = -%lf" % (i, (REF * 2 - ADC_Value[i] * REF / 0x800000)))
                        logging.info("     AVG b : ADC2 IN%d = -%lf" % (i, avg_b))
                    else:
                        val_b = ADC_Value[i] * REF / 0x7fffff
                        avg_b = moving_average_b.next(val_b)
                        logging.info("ADC2 IN%d = %lf" % (i, (ADC_Value[i] * REF / 0x7fffff)))  # 24bit
                        logging.info("     AVG b : ADC2 IN%d = %lf" % (i, avg_b))
                    
                    angle_b = (360.0 * avg_b) / REF  # degree of sea

                    current_pitch_right_buffer.append(angle_b)
                    logging.info("SEA  ID : {}   - ANGLE : ".format(i)+ format(angle, '.3f'))

                    time.sleep(1 / frequency)

            except IOError as e:
                logging.error(e)

            except KeyboardInterrupt:
                logging.info("ctrl + c: Program end")
                # logging.info(current_yaw_right_buffer,current_pitch_right_buffer)
                logging.info("sea_serial NOT in DefaultSEAChannelSer: threadName.exit() !")
                ADC.ADS1263_Exit()
                threadName.exit()
                exit()
        
        # write the sea data
        file_sea_angle = "file_sea_angle_records_"+sea_serial+".txt"
        with open(file_sea_angle, 'w') as f:
            if sea_serial == 0 and current_pitch_left_buffer != []:
                f.write("time_stamped"+"\t"+"left_shoulder_pitch_angle      "+"\n") 
                for frame in current_pitch_left_buffer:
                    # data file write as txt
                    f.write("{}   {}  ".format(time.ctime(time.time()), frame + '\r'))
            if sea_serial == 1 and current_pitch_right_buffer  != []:
                f.write("time_stamped"+"\t"+"right_shoulder_pitch_angle      "+"\n") 
                for frame in current_pitch_right_buffer:
                    # data file write as txt
                    f.write("{}   {}  ".format(time.ctime(time.time()), frame + '\r'))
    
    else:
        logging.error("%s: %s, ERROR : sea_serial NOT in DefaultSEAChannelSer !" 
                        % (threadName, time.ctime(time.time())))
        logging.info("sea_serial NOT in DefaultSEAChannelSer: threadName.exit() !")
        threadName.exit()
        


# 增加关节舵机角度的vel和iterval参数动态调整接口API
def generate_pose_vtc(angle, servo_id, vel=100, iterval=5000, t_acc=35, t_dec=35, power=10000, mean_dps=70.0):
    pose_tmp = {'angle': angle, 'is_mturn': False, 'interval': iterval,
                'vel': vel, 't_acc': t_acc,
                't_dec': t_dec, 'power': power,
                'mean_dps': mean_dps}
    gen_pose = Pose(POSE=pose_tmp, SERVO_ID=servo_id)
    return gen_pose


class MovingAverage:
    def __init__(self, size: int):
        self.que = queue.Queue(size)

    def next(self, val: float):
        if (self.que.full()):
            self.que.get()  # get the queue item value
        self.que.put(val)
        return sum(self.que.queue) / self.que.qsize()  # calculate the average voltage value


def main():
    # main function to construct new control and records threading

    # dual_arm class obj init
    dual_arm = ArmMotionControl(arm_name="dual_arm", dof=6, CONTROL_MODE="DAMPING", with_wrist=False, is_debug=True,
                                damping_power=1000)
    wrist_gripper = WristGripperControl(servo_channel=ending_servo_idx, is_debug=True)

    # 初始化舵机管理器类对象、串口端口地址
    dual_arm.creat_uservo_obj(SERVO_BAUDRATE=115200)
    if dual_arm.servo_buffer(is_scan_servo=True):
        dual_arm.logging_arm_info()
    # dual arm and gripper init
    dual_arm.arm_state_init(force_angle_init=False)
    dual_arm.set_all_servo_mode(CONTROL_MODE="CIRCLE")
    wrist_gripper.all_ending_init(motion_group=4)
    time.sleep(2)

    # 创建新线程
    controlThreadName = "control"
    recordThreadName = "record"
    seaRcd0ThreadName = "record_sea_0"
    seaRcd1ThreadName = "record_sea_1"

    # Multi Threading Thread ID [1,2,3,4]
    thread1 = dfrThread(threadID=1,
                        name=controlThreadName,
                        controlPeriod=CONTROL_PERIOD.get("SEC_3"),
                        samplingFrequency=SAMPLING_FREQUENCY,
                        dual_arm_obj=dual_arm)

    thread2 = dfrThread(threadID=2,
                        name=recordThreadName,
                        controlPeriod=None,
                        samplingFrequency=SAMPLING_FREQUENCY,
                        dual_arm_obj=dual_arm)

    thread_sea_0 = dfrThread(threadID = 3,
                        name=seaRcd0ThreadName,
                        controlPeriod=None,
                        samplingFrequency=SAMPLING_FREQUENCY,
                        dual_arm_obj=None,
                        sea_serial= 0 )

    thread_sea_1 = dfrThread(threadID = 4,
                        name=seaRcd1ThreadName,
                        controlPeriod=None,
                        samplingFrequency=SAMPLING_FREQUENCY,
                        dual_arm_obj=None,
                        sea_serial=1)

    # 开启新线程
    thread1.start()
    thread2.start()
    thread_sea_0.start()
    thread_sea_1.start()

    # 等待线程完成
    thread1.join()
    thread2.join()
    thread_sea_0.join()
    thread_sea_1.join()

    logging.info("退出主线程")


# main function
if __name__ == '__main__':
    # 创建一个logger_debug
    import sys

    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)  # Log等级总开关-debug
    print("Logging Set Info: logger.setLevel(logging.DEBUG)  # Log等级总开关-debug")
    try:
        main()
    except KeyboardInterrupt:
        print("Shutting down robot arm debug.")